'''
Brief：
  CNN训练相关代码，获取数据集
author：GL
date：2021/01
'''

import torch
import torchvision
import cv2
import numpy as np
class ImageRecord(object):
    '''
    Brief:
      数据集的每一个record
      一条记录的格式:path(str),label(int)
    '''
    def __init__(self, row):
        self._data = row

    @property            
    def path(self):
        """
        Brief:
          记录的图像文件路径，建议使用绝对路径
          @是装饰器用法，@property装饰器可以使将函数用类属性的方式调用
          这里不需要管
        Return:
          seld._data[0](str):图像文件路径
        """
        return self._data[0]
    @property
    def label(self):
        return int(self._data[1])

class Dataset(object):
    """
    加载灰度图像的dataset
    """
    def __init__(self,list_file):
        '''
        Args:
          list_file(str):训练集文件,由gen_dataset_list生成
        '''
        self.list_file = list_file
        self._parse_list()

    def _parse_list(self):
        """
        Brief:
          解析初始化时传入的数据集文件
        """
        self.img_list = [ImageRecord(x.strip().split(' ')) for x in open(self.list_file)]
        #这是生成器(generator)的用法，暂时不需要管，等价于
        #self.img_list = []
        #for x in open(self.list_file):
        #    temp = ImageRecord(x.strip().split(' ')) 
        #    self.img_list.append(temp)

    def get(self, record):
        '''
        Brief:处理图像并返回一个处理完成的record
        Args:
          record(ImageRecord):ImageRecord对象
        Return:
          
        '''

        #---------------------实现读取图像并转换为torch.Tensor代码------------------------
        #要求：使用opencv-python读取灰度图像，转换为torch.Tensor,类型为float32
        #预处理完成后Tensor的尺寸自己定义，能传入自己的网络即可
        src = cv2.imread(record.path,0)#读取黑白图
        src = cv2.resize(src,(52,52))
        image = np.expand_dims(src,axis=2)# 增加维度
        process_data  =torch.from_numpy(image).permute(2, 0, 1).contiguous()
        return process_data, record.label

    def __getitem__(self, index):
        """
        Brief:
          特殊方法，torch在调用DataLoader时就是调用这个方法
        """
        record = self.img_list[index]
        return self.get(record)

    def __len__(self):
        return len(self.img_list)

    
